use std::fs::File;
use std::io::{BufRead, BufReader, Error, Seek, SeekFrom};
use std::path::Path;

use std::sync::{Arc, RwLock};
use std::thread::spawn;

use rand::distributions::{WeightedError, WeightedIndex};
use rand::prelude::*;

fn main() {
    let n2_dir = "D:/PythonProject/ngram/ngram2";
    let cache_n2 = make_cache(n2_dir).unwrap();
    let arc_cache_n2 = Arc::new(RwLock::new(cache_n2));
    let n3_dir = "D:/PythonProject/ngram/ngram3";
    let cache_n3 = make_cache(n3_dir).unwrap();
    let arc_cache_n3 = Arc::new(RwLock::new(cache_n3));
    let mut generate_list = vec!["结果".to_owned(),"应当".to_owned()];

    let thread_counts = 12;
    for i in 0..20 {
        let instant = std::time::Instant::now();
        let (next, score) = speed_up_get_next_word_n2(
            arc_cache_n2.clone(),
            generate_list.last().unwrap().clone(),
            thread_counts,
        );

        print!("{:?} ", instant.elapsed());
        if next.len()!=0 {
            generate_list.push(next);
        }
        println!("{} ", generate_list.concat());

        let len = generate_list.len();
        assert!(len >= 2);
        let instant = std::time::Instant::now();
        let (next, _) = speed_up_get_next_word_n3(
            arc_cache_n3.clone(),
            generate_list[len - 2..len].concat().to_owned(),
            thread_counts,
        );
        print!("{:?} ", instant.elapsed());
        if next.len()!=0 {
            generate_list.push(next);
        }
        println!("{} ", generate_list.concat());
    }
}

fn speed_up_get_next_word_n2(
    file_cache: Arc<RwLock<Vec<File>>>,
    search_string: String,
    thread_counts: usize,
) -> (String, usize) {
    let mut threads = vec![];
    for i in 0..thread_counts {
        let search = search_string.clone();
        let arc = file_cache.clone();
        let handle = spawn(move || {
            let locker = arc.read().unwrap();
            let per_thread_cache = locker.len() / thread_counts;
            let next_word = get_next_word_n2(
                &locker[i * per_thread_cache..(i + 1) * per_thread_cache],
                &search,
            );
            return next_word;
        });
        threads.push(handle);
    }

    let mut data = vec![];
    for t in threads {
        if let Ok(r) = t.join() {
            data.push(r)
        }
    }
    let sample = uniform_sample(data);

    let mut locker = file_cache.write().unwrap();
    reset_cache(&mut locker);
    sample.unwrap_or_else(|_| ("".to_owned(), 0))
}

fn speed_up_get_next_word_n3(
    file_cache: Arc<RwLock<Vec<File>>>,
    search_string: String,
    thread_counts: usize,
) -> (String, usize) {
    let mut threads = vec![];
    for i in 0..thread_counts {
        let search = search_string.clone();
        let arc = file_cache.clone();
        let handle = spawn(move || {
            let locker = arc.read().unwrap();
            let per_thread_cache = locker.len() / thread_counts;
            let next_word = get_next_word_n3(
                &locker[i * per_thread_cache..(i + 1) * per_thread_cache],
                &search,
            );
            return next_word;
        });
        threads.push(handle);
    }

    let mut data = vec![];
    for t in threads {
        if let Ok(r) = t.join() {
            data.push(r)
        }
    }
    let sample = uniform_sample(data);

    let mut locker = file_cache.write().unwrap();
    reset_cache(&mut locker);
    sample.unwrap_or_else(|_| ("".to_owned(), 0))

}
fn get_next_word_n2(cache: &[File], search: &str) -> (String, usize) {
    let mut ret = vec![];
    for file in cache {
        let lines = search_line_start_with_n2(file, search);
        for line in lines {
            let split = line.split_whitespace().collect::<Vec<_>>();
            ret.push((split[1].to_owned(), split[3].parse::<usize>().unwrap()));
        }
    }
    max_sample(ret)
}
fn get_next_word_n3(cache: &[File], search: &str) -> (String, usize) {
    let mut ret = vec![];
    for file in cache {
        let lines = search_line_start_with_n3(file, search);
        for line in lines {
            let split = line.split_whitespace().collect::<Vec<_>>();
            ret.push((split[2].to_owned(), split[4].parse::<usize>().unwrap()));
        }
    }
    max_sample(ret)
}
fn uniform_sample(data: Vec<(String, usize)>) -> Result<(String, usize), WeightedError> {
    let dist = WeightedIndex::new(data.iter().map(|item| item.1))?;
    let mut rng = thread_rng();
    let index: usize = dist.sample(&mut rng);
    Ok(data[index].clone())
}

fn max_sample(data: Vec<(String, usize)>) -> (String, usize) {
    let mut ret = (String::new(), 0);
    for d in data {
        if d.1 > ret.1 {
            ret = d;
        }
    }
    ret
}
fn search_line_start_with_n2(file: &File, search: &str) -> Vec<String> {
    let mut reader = BufReader::new(file);
    let mut ret = vec![];
    loop {
        let mut buffer = String::new();

        match reader.read_line(&mut buffer) {
            Ok(s) if s == 0 => break,
            Ok(s) if buffer.starts_with(search) => ret.push(buffer.clone()),
            Err(_) => {}
            _ => {}
        }
    }
    ret
}
fn starts_with_ignore_whitespace(text: &str, prefix: &str) -> bool {
    let text_iter = text.chars().filter(|c| !c.is_whitespace());
    let prefix_iter = prefix.chars().filter(|c| !c.is_whitespace());

    for (t, p) in text_iter.zip(prefix_iter) {
        if t != p {
            return false;
        }
    }

    // 确保prefix的所有字符都被匹配
    true
}

fn search_line_start_with_n3(file: &File, search: &str) -> Vec<String> {
    let mut reader = BufReader::new(file);
    let mut ret = vec![];
    loop {
        let mut buffer = String::new();

        match reader.read_line(&mut buffer) {
            Ok(s) if s == 0 => break,
            Ok(s) if starts_with_ignore_whitespace(&buffer, search) => ret.push(buffer.clone()),
            Err(_) => {}
            _ => {}
        }
    }
    ret
}
fn make_cache(dir_path: &str) -> Result<Vec<File>, Error> {
    let dir = Path::new(dir_path);
    let mut cache = Vec::with_capacity(100);
    if dir.is_dir() {
        for entry in std::fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_file() {
                println!("{:?}", path.display());
                cache.push(File::open(path)?);
            }
        }
    }
    Ok(cache)
}

fn reset_cache(cache: &mut [File]) {
    for file in cache {
        file.seek(SeekFrom::Start(0)).unwrap();
    }
}
